// Copyright (c) The Diem Core Contributors
// Copyright (c) The Move Contributors
// SPDX-License-Identifier: Apache-2.0

use std::{
    collections::BTreeMap,
    io::{BufRead, Write},
    path::{Path, PathBuf},
};

use clap::ArgAction;
use clap::Parser;
use serde::{Deserialize, Serialize};

use move_compiler::{
    editions::{Edition, Flavor},
    shared::NumericalAddress,
};
use move_core_types::{account_address::AccountAddress, identifier::Identifier};
use move_model_2::source_model;
use move_package_alt::{
    graph::NamedAddress,
    schema::{EnvironmentName, ModeName, OriginalID},
};
use move_symbol_pool::Symbol;

use move_package_alt::{flavor::MoveFlavor, package::RootPackage, schema::Environment};

use crate::{
    build_plan::BuildPlan,
    compiled_package::{BuildNamedAddresses, CompiledPackage},
    migrate::migrate,
    model_builder,
};

use super::lint_flag::LintFlag;

#[derive(Debug, Parser, Clone, Serialize, Deserialize, Eq, PartialEq, PartialOrd, Default)]
#[clap(about)]
pub struct BuildConfig {
    /// Compile in 'test' mode and include the code in the 'tests' directory.
    #[clap(name = "test-mode", long = "test", global = true)]
    pub test_mode: bool,

    /// Generate documentation for packages
    #[clap(name = "generate-docs", long = "doc", global = true)]
    pub generate_docs: bool,

    /// Save disassembly for generated bytecode along with
    /// bytecode maps (source maps for disassembled bytecode)
    #[clap(name = "save-disassembly", long = "disassemble", global = true)]
    pub save_disassembly: bool,

    /// Installation directory for compiled artifacts. Defaults to current directory.
    #[clap(long = "install-dir", global = true)]
    pub install_dir: Option<PathBuf>,

    /// Force recompilation of all packages
    #[clap(name = "force-recompilation", long = "force", global = true)]
    pub force_recompilation: bool,

    /// Default flavor for move compilation, if not specified in the package's config
    #[clap(long = "default-move-flavor", global = true)]
    pub default_flavor: Option<Flavor>,

    /// Default edition for move compilation, if not specified in the package's config
    #[clap(long = "default-move-edition", global = true)]
    pub default_edition: Option<Edition>,

    /// Optional location to save the lock file to, if package resolution succeeds.
    #[clap(skip)]
    pub lock_file: Option<PathBuf>,

    /// If set, ignore any compiler warnings
    #[clap(long = move_compiler::command_line::SILENCE_WARNINGS, global = true)]
    pub silence_warnings: bool,

    /// If set, warnings become errors
    #[clap(long = move_compiler::command_line::WARNINGS_ARE_ERRORS, global = true)]
    pub warnings_are_errors: bool,

    /// If set, reports errors at JSON
    #[clap(long = move_compiler::command_line::JSON_ERRORS, global = true)]
    pub json_errors: bool,

    /// Additional named address mapping. Useful for tools in rust
    #[clap(skip)]
    pub additional_named_addresses: BTreeMap<String, AccountAddress>,

    #[clap(flatten)]
    pub lint_flag: LintFlag,

    /// Arbitrary mode -- this will be used to enable or filter user-defined `#[mode(<MODE>)]`
    /// annotations during compilation.
    #[arg(
        long = "mode",
        value_name = "MODE",
        value_parser = parse_symbol,
        action = ArgAction::Append,
        global = true
    )]
    pub modes: Vec<Symbol>,

    /// Forces use of lock file without checking if it needs to be updated
    /// (regenerates it only if it doesn't exist)
    #[clap(skip)]
    pub force_lock_file: bool,

    /// Forces the `root` package to have `0x0` as its address, instead of the published
    /// address (if it exists). Useful for `upgrade` operations (or any other cases we might want to..)
    #[clap(skip)]
    pub root_as_zero: bool,

    #[clap(
        long = "environment",
        short = 'e',
        global = true,
        help = "Environment to use for building packages"
    )]
    pub environment: Option<EnvironmentName>,

    /// If set, any dependencies that are not published will have their address set to 0x0.
    #[clap(skip)]
    pub set_unpublished_deps_to_zero: bool,
}

impl BuildConfig {
    pub async fn compile_package<F: MoveFlavor, W: Write + Send>(
        &self,
        path: &Path,
        env: &Environment,
        writer: &mut W,
    ) -> anyhow::Result<CompiledPackage> {
        let root_pkg = RootPackage::<F>::load(path, env.clone(), self.mode_set()).await?;
        BuildPlan::create(&root_pkg, self)?.compile(writer, |compiler| compiler)
    }

    /// Migrate the package at `path`.
    pub async fn migrate_package<F: MoveFlavor, W: Write + Send, R: BufRead>(
        mut self,
        path: &Path,
        env: Environment,
        writer: &mut W,
        reader: &mut R,
    ) -> anyhow::Result<()> {
        // we set test to migrate all the code
        self.test_mode = true;
        let root_pkg = RootPackage::<F>::load(path, env, self.mode_set()).await?;
        let build_plan = BuildPlan::create(&root_pkg, &self)?;

        migrate(build_plan, writer, reader)?;
        Ok(())
    }

    pub async fn move_model_from_path<F: MoveFlavor, W: Write + Send>(
        &self,
        path: &Path,
        env: Environment,
        writer: &mut W,
    ) -> anyhow::Result<source_model::Model> {
        let root_pkg = RootPackage::<F>::load(path, env, self.mode_set()).await?;
        self.move_model_from_root_pkg(&root_pkg, writer).await
    }

    pub async fn move_model_from_root_pkg<F: MoveFlavor, W: Write + Send>(
        &self,
        root_pkg: &RootPackage<F>,
        writer: &mut W,
    ) -> anyhow::Result<source_model::Model> {
        model_builder::build(writer, root_pkg, self)
    }

    /// Build the addresse for the supplied config, so we can inject 0x0s etc.
    pub fn addresses_for_config(
        &self,
        named_addresses: BTreeMap<Identifier, NamedAddress>,
    ) -> BuildNamedAddresses {
        // Make unpublished deps `0x0` if the config says so.
        let fixed_named_addresses = named_addresses
            .into_iter()
            .map(|(id, addr)| {
                (
                    id,
                    match addr {
                        NamedAddress::Unpublished { dummy_addr: _ }
                            if self.set_unpublished_deps_to_zero =>
                        {
                            NamedAddress::Defined(OriginalID(AccountAddress::ZERO))
                        }
                        addr => addr,
                    },
                )
            })
            .collect();

        // Force root as zero for upgrade operations.
        let mut addresses: BuildNamedAddresses = if self.root_as_zero {
            BuildNamedAddresses::root_as_zero(fixed_named_addresses)
        } else {
            fixed_named_addresses.into()
        };

        // Inject additional named addresses.
        for (pkg, address) in self.additional_named_addresses.clone() {
            addresses.inner.insert(
                pkg.clone().into(),
                NumericalAddress::new(
                    address.into_bytes(),
                    move_compiler::shared::NumberFormat::Hex,
                ),
            );
        }

        addresses
    }

    /// Produce the set of mode names to hand to the package system.
    pub fn mode_set(&self) -> Vec<ModeName> {
        let mut result: Vec<ModeName> = self.modes.iter().map(|mode| mode.to_string()).collect();

        if self.test_mode {
            result.push("test".to_string());
        }

        result
    }
}

fn parse_symbol(s: &str) -> Result<Symbol, String> {
    Ok(Symbol::from(s))
}
